import torch
from torch.nn import Module
from torch.nn import functional as F


class LabelSmoothingLoss(Module):

    def __init__(self, num_classes, eps=0.1, reduction="mean", **kwargs):
        super(LabelSmoothingLoss, self).__init__(**kwargs)
        self.eps = eps
        self.num_classes = num_classes
        self.K = eps / (num_classes - 1)
        self.reduction = reduction

    def forward(self, input, target):
        softmax_input = F.log_softmax(input, dim=1)
        onehot = input.new_ones(target.size(0), self.num_classes) * self.K
        onehot.scatter_(dim=1, index=target.unsqueeze(1), value=1 - self.eps)
        loss = - torch.sum(onehot * softmax_input, dim=-1)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        elif self.reduction == 'none':
            return loss
